import pygame
import sys
import random

# 初始化pygame
pygame.init()

# 游戏配置
SCREEN_WIDTH, SCREEN_HEIGHT = 640, 480
CELL_SIZE = 20
CELL_NUMBER_X = SCREEN_WIDTH // CELL_SIZE
CELL_NUMBER_Y = SCREEN_HEIGHT // CELL_SIZE
FPS = 10

# 颜色定义
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
WHITE = (255, 255, 255)

# 设置屏幕
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption('贪吃蛇')
clock = pygame.time.Clock()

# 字体
font = pygame.font.SysFont(None, 36)

class Snake:
    def __init__(self):
        # 蛇的初始位置 (x, y)
        self.body = [pygame.Vector2(5, 10), pygame.Vector2(4, 10), pygame.Vector2(3, 10)]
        self.direction = pygame.Vector2(1, 0)  # 初始方向向右
        self.new_block = False

    def draw_snake(self):
        for block in self.body:
            x_pos = int(block.x * CELL_SIZE)
            y_pos = int(block.y * CELL_SIZE)
            block_rect = pygame.Rect(x_pos, y_pos, CELL_SIZE, CELL_SIZE)
            pygame.draw.rect(screen, GREEN, block_rect)

    def move_snake(self):
        if self.new_block:
            body_copy = self.body[:]
            body_copy.insert(0, body_copy[0] + self.direction)
            self.body = body_copy[:]
            self.new_block = False
        else:
            body_copy = self.body[:-1]
            body_copy.insert(0, body_copy[0] + self.direction)
            self.body = body_copy[:]

    def add_block(self):
        self.new_block = True

    def check_collision(self):
        # 检查是否撞墙
        if not 0 <= self.body[0].x < CELL_NUMBER_X or not 0 <= self.body[0].y < CELL_NUMBER_Y:
            return True
        # 检查是否撞到自己
        for block in self.body[1:]:
            if block == self.body[0]:
                return True
        return False

class Food:
    def __init__(self):
        self.randomize()

    def draw_food(self):
        food_rect = pygame.Rect(int(self.pos.x * CELL_SIZE), int(self.pos.y * CELL_SIZE), CELL_SIZE, CELL_SIZE)
        pygame.draw.rect(screen, RED, food_rect)

    def randomize(self):
        self.x = random.randint(0, CELL_NUMBER_X - 1)
        self.y = random.randint(0, CELL_NUMBER_Y - 1)
        self.pos = pygame.Vector2(self.x, self.y)

def main():
    snake = Snake()
    food = Food()
    score = 0

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    if snake.direction.y != 1:
                        snake.direction = pygame.Vector2(0, -1)
                if event.key == pygame.K_DOWN:
                    if snake.direction.y != -1:
                        snake.direction = pygame.Vector2(0, 1)
                if event.key == pygame.K_RIGHT:
                    if snake.direction.x != -1:
                        snake.direction = pygame.Vector2(1, 0)
                if event.key == pygame.K_LEFT:
                    if snake.direction.x != 1:
                        snake.direction = pygame.Vector2(-1, 0)

        snake.move_snake()

        # 检查是否吃到食物
        if snake.body[0] == food.pos:
            food.randomize()
            snake.add_block()
            score += 1

        # 检查碰撞
        if snake.check_collision():
            game_over_screen(score)
            return

        screen.fill(BLACK)
        food.draw_food()
        snake.draw_snake()
        
        # 显示分数
        score_text = font.render(f'Score: {score}', True, WHITE)
        screen.blit(score_text, (10, 10))
        
        pygame.display.update()
        clock.tick(FPS)

def game_over_screen(score):
    screen.fill(BLACK)
    game_over_text = font.render('Game Over!', True, WHITE)
    score_text = font.render(f'Final Score: {score}', True, WHITE)
    restart_text = font.render('Press R to Restart or Q to Quit', True, WHITE)
    
    screen.blit(game_over_text, (SCREEN_WIDTH//2 - game_over_text.get_width()//2, SCREEN_HEIGHT//2 - 60))
    screen.blit(score_text, (SCREEN_WIDTH//2 - score_text.get_width()//2, SCREEN_HEIGHT//2 - 20))
    screen.blit(restart_text, (SCREEN_WIDTH//2 - restart_text.get_width()//2, SCREEN_HEIGHT//2 + 20))
    
    pygame.display.update()
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    main()  # 重新开始游戏
                if event.key == pygame.K_q:
                    pygame.quit()
                    sys.exit()

if __name__ == '__main__':
    main()